import asyncio
import json
from contextlib import asynccontextmanager

import aioredis
from typing import Dict, List
import uvicorn
from loguru import logger
from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect, HTTPException, Request, Form, Body
from fastapi.middleware.cors import CORSMiddleware
from config import WebSocketToken


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时的逻辑
    print("Starting up...")
    redis_task = asyncio.create_task(listen_to_redis())
    try:
        yield
    finally:
        # 关闭时的逻辑
        print("Shutting down...")
        redis_task.cancel()
        await redis_task


app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法，包括OPTIONS
    allow_headers=["*"],
)


# 用于存储连接的 WebSocket 客户端
class ConnectionManager:
    def __init__(self):
        # 存储所有活跃的 WebSocket 连接
        self.active_connections: List[WebSocket] = []
        # 存储房间信息，每个房间对应一个 WebSocket 列表
        self.rooms: Dict[str, List[WebSocket]] = {}

    async def connect(self, websocket: WebSocket):
        # 接受WebSocket连接
        await websocket.accept()
        # 将WebSocket连接添加到活动连接列表中
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
        # 从所有房间中移除断开的 WebSocket 连接
        for room in self.rooms.values():
            if websocket in room:
                room.remove(websocket)

    async def send_personal_message(self, message, websocket: WebSocket):
        # 通过WebSocket发送消息
        await websocket.send_text(json.dumps(message, ensure_ascii=False))
    async def broadcast(self, message: str):
        # 遍历活动连接
        for connection in self.active_connections:
            # 向连接发送文本消息
            await connection.send_text(message)

    # 广播到指定房间
    async def broadcast_to_room(self, room_id, message) -> dict:
        results = {"success": True, "details": []}
        logger.info(f"当前房间:{type(room_id),room_id},  {self.rooms}")
        if str(room_id) in self.rooms:
            for connection in self.rooms[str(room_id)]:
                try:
                    await connection.send_text(json.dumps(message, ensure_ascii=False))
                    logger.info(f"消息已发送到连接: {connection}")
                    results["details"].append({"connection": connection, "status": "success"})
                except Exception as e:
                    logger.error(f"发送消息到连接失败: {connection}, 错误: {e}")
                    results["details"].append({"connection": connection, "status": "failed", "error": str(e)})
                    results["success"] = False
        else:
            logger.warning(f"房间未找到: {room_id}")
            results["success"] = False
            results["details"].append({"error": "Room not found"})
        return results

    def join_room(self, room_id: str, websocket: WebSocket):
        if room_id not in self.rooms:
            self.rooms[room_id] = []
        if websocket not in self.rooms[room_id]:
            self.rooms[room_id].append(websocket)
        logger.info(f"WebSocket连接加入房间: {room_id}, 当前房间连接数: {len(self.rooms[room_id])}")

    def leave_room(self, room_id: str, websocket: WebSocket):
        if room_id in self.rooms and websocket in self.rooms[room_id]:
            self.rooms[room_id].remove(websocket)
            logger.info(f"WebSocket连接离开房间: {room_id}, 当前房间连接数: {len(self.rooms[room_id])}")
            if not self.rooms[room_id]:
                del self.rooms[room_id]
                logger.info(f"房间已清空并删除: {room_id}")

# 监听redis消息
async def listen_to_redis():
    # 创建 Redis 连接
    redis = aioredis.from_url('redis://localhost')
    # 创建 Redis 发布订阅对象
    pubsub = redis.pubsub()
    # 订阅 'websocket_channel' 频道
    await pubsub.subscribe('websocket_channel')
    try:
        # 异步监听 Redis 发布订阅消息
        async for message in pubsub.listen():
            # 判断消息类型是否为消息
            if message['type'] == 'message':
                # 获取消息内容
                msg = message['data']
                # 解析 JSON 格式的消息内容
                msg = json.loads(msg)
                # 获取房间 ID
                room_id = msg.get("room_id")
                # 获取用例状态
                case_status = msg.get("case_status")
                # 如果用例状态存在
                if case_status:
                    # 打印收到的房间信息和用例状态
                    print(f"收到房间信息 {room_id}: {case_status}")
                    # 向指定房间广播用例状态
                    await manager.broadcast_to_room(str(room_id), case_status)
    except asyncio.CancelledError:
        # 捕获任务取消异常
        print("Redis 侦听器任务已取消.")
    finally:
        # 取消订阅 'websocket_channel' 频道
        await pubsub.unsubscribe('websocket_channel')
        # 关闭 Redis 连接
        await redis.close()


manager = ConnectionManager()


@app.websocket("/api/ws/{room_id}")
async def websocket_endpoint(websocket: WebSocket, room_id, token: str):
    if token != WebSocketToken:
        raise HTTPException(status_code=403, detail="token错误")
    await manager.connect(websocket)
    manager.join_room(room_id, websocket)
    try:
        await manager.send_personal_message({"message": f"已连接到服务器"}, websocket)
        while True:
            data = await websocket.receive_text()
            logger.info(f"接收到消息: {data}")
            # 如果需要将字符串解析为 JSON 对象
            try:
                json_data = json.loads(data)  # 将字符串转换为字典
                broadcast_result = await manager.broadcast_to_room(str(room_id), json_data)
                logger.info(f"转发结果: {broadcast_result}")
            except json.JSONDecodeError:
                print("接收到的消息不是有效的 JSON 格式")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        manager.leave_room(room_id, websocket)
        print("客户端已断开连接")


if __name__ == '__main__':
    uvicorn.run(
        "websocket_app:app",
        host="127.0.0.1",
        port=5005,
        workers=1,
        reload=False,
    )
